package cn.tworice.cryption.util;


import javax.crypto.Cipher;
import javax.crypto.SecretKey;
import javax.crypto.spec.SecretKeySpec;
import java.util.Base64;
/**
 * 属性基加密工具类
 */
public class AttributeBasedEncryption {

    /**
     * 根据角色构建密钥
     * @param role 角色
     * @return 密钥字符串
     */
    public static String getKeyString(String role) {
        String keyString = Base64.getEncoder().encodeToString(role.getBytes());;
        boolean flag = true;
        while (flag){
            keyString += keyString;
            if (keyString.length() > 16) {
                flag = false;
            }
        }
        return keyString.substring(0, 16);
    }

    /**
     * 加密方法
     * @param data 明文
     * @param key 密钥
     * @return 密文
     */
    public static String encrypt(String data, SecretKey key) throws Exception {
        Cipher cipher = Cipher.getInstance("AES");
        cipher.init(Cipher.ENCRYPT_MODE, key);
        byte[] encryptedData = cipher.doFinal(data.getBytes());
        return Base64.getEncoder().encodeToString(encryptedData);
    }

    /**
     * 解密方法
     * @param encryptedData 加密数据
     * @param key 密钥
     * @return 铭文
     */
    public static String decrypt(String encryptedData, SecretKey key) throws Exception {
        Cipher cipher = Cipher.getInstance("AES");
        cipher.init(Cipher.DECRYPT_MODE, key);
        byte[] decodedData = Base64.getDecoder().decode(encryptedData);
        byte[] decryptedData = cipher.doFinal(decodedData);
        return new String(decryptedData);
    }

    /**
     * 根据角色加密数据
     * @param data 数据
     * @param role 角色
     * @return 密文
     */
    public static String encryptWithRole(String data, String role) throws Exception {
        // 基于角色生成不同的密钥
        // 在实际的ABE系统中，角色会影响加密密钥的生成
        return encrypt(data, new SecretKeySpec(getKeyString(role).getBytes(), "AES"));
    }

    /**
     * 根据角色解密数据
     * @param encryptedData 密文
     * @param role 角色
     * @return 明文
     */
    public static String decryptWithRole(String encryptedData, String role) throws Exception {
        // 基于角色生成不同的密钥
        String keyString = getKeyString(role);
        System.out.println("解密: " + keyString);
        SecretKey key = new SecretKeySpec(keyString.getBytes(), "AES");
        // 解密
        return decrypt(encryptedData, key);
    }

    public static void main(String[] args) throws Exception {
        // 数据加密和解密示例
        String data = "This is a secret message";
        String role = "管理员";  // 假设用户角色为admin

        // 加密
        String encryptedData = encryptWithRole(data, role);
        System.out.println("Encrypted Data: " + encryptedData);

        // 解密
        String decryptedData = decryptWithRole(encryptedData, role);
        System.out.println("Decrypted Data: " + decryptedData);
    }
}

